from motor import motor_asyncio
from bson.objectid import ObjectId
from bson.int64 import Int64


class MongoDB(object):
    def __init__(self, config):
        self.config = config
        self.mongodb = None
        self.__getConn()

    # 链接数据库
    def __getConn(self):
        if self.mongodb:
            return self.mongodb

        _kwargs = {
            "maxPoolSize": self.config['maxPoolSize'],
            "connect": False,
        }
        self.client = motor_asyncio.AsyncIOMotorClient(self.config["url"], **_kwargs)
        self.mongodb = self.client.get_database(self.config["dbname"])

        return self.mongodb

    # 转化为objectId数据类型
    def toObjectId(self, oid):
        return ObjectId(oid)

    # 转化为int64数据类型
    def toInt64(self, num):
        return Int64(num)

    # 原始聚合
    async def aggregate(self, tb, pipeline):
        _res = await self.mongodb[tb].aggregate(pipeline)
        return _res

    # 插入数据
    async def insert_many(self, tb, *arg, **key):
        res = await self.mongodb[tb].insert_many(*arg, **key)
        return res

    # 插入数据
    async def insert_one(self, tb, *arg, **key):
        key["bypass_document_validation"] = True
        res = await self.mongodb[tb].insert_one(*arg, **key)
        return res

    # 修改数据
    async def update(self, tb, where, data, **key):
        n_data = data
        # 存在其他$使用，只能用$set
        if (''.join(data.keys()).find('$')) == -1:
            n_data = {"$set": data}

        res = await self.mongodb[tb].update_many(where, n_data, **key)
        return res

    @staticmethod
    def filter_key(key):
        if "where" in key:
            key['filter'] = key['where']
            del key['where']

        if "fields" in key and type(key["fields"]) == list:
            f_map = {}
            for ele in key["fields"]:
                if ele != "_id":
                    f_map[ele] = 1
                else:
                    f_map[ele] = 0

            key["fields"] = f_map

        if "fields" in key:
            key["projection"] = key["fields"]
            del key["fields"]

    async def find1(self, tb, *arg, **key):
        self.filter_key(key)
        res = await self.mongodb[tb].find_one(*arg, **key)
        return res

    # 查询数据集
    async def find(self, tb, *arg, **key):
        self.filter_key(key)
        res = await self.mongodb[tb].find(*arg, **key).to_list(None)
        return res

    # 删除数据
    async def delete(self, tb, *arg, **key):
        if "where" in key:
            key["spec_or_id"] = key["where"]
            del key['where']
        res = await self.mongodb[tb].delete_many(*arg, **key)
        return res

    # 计数命令
    async def count(self, tb, where=None, **kwargs):
        res = await self.mongodb[tb].count_documents(where, **kwargs)
        return res


if __name__ == "__main__":
    pass
